package cn.linkai.star.graph.comfyui.client.config;

import cn.linkai.star.graph.comfyui.client.api.ComfyuiApi;
import cn.linkai.star.graph.comfyui.client.handler.ComfyuiMessageHandler;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.WebSocketConnectionManager;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.time.Duration;

@Configuration
public class ComfyuiConfig {

    @Bean
    public ComfyuiApi comfyuiApi() {
        //创建日志拦截器
        HttpLoggingInterceptor loggingInterceptor = new HttpLoggingInterceptor ();
        /**
         *   - NONE：不输出
         *   - BASIC：输出基本摘要
         *   - HEADERS：输出头信息
         *   - BODY：输出body数据
         */
        //配置日志级别
        loggingInterceptor.setLevel (HttpLoggingInterceptor.Level.BODY);
        //创建okHttpClient对象
        OkHttpClient okHttpClient = new OkHttpClient.Builder ()
                //添加日志拦截器
                .addInterceptor (loggingInterceptor)
                //设置连接超时时间
                .connectTimeout (Duration.ofSeconds (30))
                //设置读取超时时间
                .readTimeout (Duration.ofSeconds (30))
                //设置重连
                .retryOnConnectionFailure (true)
                .build ();
        //创建 Retrofit 对象
        Retrofit retrofit = new Retrofit.Builder ()
                //指定comfyui的地址
                .baseUrl ("http://127.0.0.1:8188")
                //配置retrofit客户端okHttp
                .client (okHttpClient)
                //指定反序化数据的处理方式-jackson
                .addConverterFactory (JacksonConverterFactory.create ())
                .build ();
        //基于Retrofit创建接口代理对象
        return retrofit.create (ComfyuiApi.class);
    }

    /**
     * 消息处理器类
     * @return
     */
    @Bean
    public ComfyuiMessageHandler comfyuiMessageHandler(){
        return new ComfyuiMessageHandler ();
    }

    /**
     * webSocket连接管理器
     * @return
     */
    @Bean
    public WebSocketConnectionManager webSocketConnectionManager(ComfyuiMessageHandler comfyuiMessageHandler){
        //定义webSocket客户端
        WebSocketClient client = new StandardWebSocketClient ();
        //定义webSocket连接地址
        String url = "ws://127.0.0.1:8188/ws?clientId=star-graph";
        //定义webSocket连接管理器
        WebSocketConnectionManager manager = new WebSocketConnectionManager (client,comfyuiMessageHandler,url);
        manager.start ();
        return manager;
    }
}
